use std::io::{Read, Write};
use std::net::TcpStream;
use std::ops::{Deref, DerefMut};

use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use pyo3_stub_gen::{define_stub_info_gatherer, derive::*};
#[gen_stub_pyclass]
#[pyclass(subclass)]
struct TcpConnector {
    url: String,
    port: u16,
    tcp_stream: Option<TcpStream>,
}

impl Deref for TcpConnector {
    type Target = TcpStream;

    fn deref(&self) -> &Self::Target {
        self.tcp_stream.as_ref().unwrap()
    }
}

impl DerefMut for TcpConnector {
    fn deref_mut(&mut self) -> &mut TcpStream {
        self.tcp_stream.as_mut().unwrap()
    }
}

#[gen_stub_pymethods]
#[pymethods]
impl TcpConnector {
    /// Formats the sum of two numbers as string.
    #[new]
    fn new(url: String, port: u16) -> PyResult<Self> {
        Ok(TcpConnector {
            tcp_stream: Some(
                TcpStream::connect(format!("{}:{}", url, port))
                    .map_err(|e| PyTypeError::new_err(e.to_string()))?,
            ),
            url,
            port,
        })
    }

    fn read(&mut self, count: usize) -> PyResult<(Vec<u8>, usize)> {
        let mut buf = vec![0; count];
        let read = self
            .deref_mut()
            .read(&mut buf)
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok((buf[0..read].to_vec(), read))
    }
    fn read_all_str(&mut self) -> PyResult<(String, usize)> {
        let mut buf = String::new();
        let read = self
            .deref_mut()
            .read_to_string(&mut buf)
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok((
            buf,
            read,
        ))
    }

    fn read_to_str(&mut self, count: usize) -> PyResult<(String, usize)> {
        let mut buf = vec![0; count];
        let read = self
            .deref_mut()
            .read(&mut buf)
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok((
            String::from_utf8(buf[0..read].to_vec())
                .map_err(|e| PyTypeError::new_err(e.to_string()))?,
            read,
        ))
    }

    fn read_to_end(&mut self) -> PyResult<(Vec<u8>, usize)> {
        let mut buf = vec![];
        let read = self
            .deref_mut()
            .read_to_end(&mut buf)
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok((buf, read))
    }

    fn write(&mut self, data: Vec<u8>) -> PyResult<usize> {
        self.deref_mut()
            .write(&data)
            .map_err(|e| PyTypeError::new_err(e.to_string()))
    }

    fn write_all(&mut self, data: Vec<u8>) -> PyResult<()> {
        self.deref_mut()
            .write_all(&data)
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok(())
    }

    fn write_str(&mut self, data: String) -> PyResult<usize> {
        let written = self
            .deref_mut()
            .write(data.as_bytes())
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok(written)
    }
    fn write_all_str(&mut self, data: String) -> PyResult<()> {
        self.deref_mut()
            .write_all(data.as_bytes())
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok(())
    }
    fn close(&mut self) -> PyResult<()> {
        let tcp_stream = self.tcp_stream.take();
        Ok(())
    }
}
/// A Python module implemented in Rust.
#[pymodule]
fn r_socket(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<TcpConnector>()?;
    Ok(())
}
define_stub_info_gatherer!(stub_info);